// Copyright (c) 2005-2008, Simon Howard
// Permission to use, copy, modify, and/or distribute this software
// for any purpose with or without fee is hereby granted, provided
// that the above copyright notice and this permission notice appear
// in all copies.
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
// CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
// NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

#![allow(dead_code)]

use std::cmp::Ordering;
use std::mem;

const NUM_TEST_VALUES: usize = 1000;

#[derive(Debug)]
struct AVLTreeNode<K, V> {
    key: K,
    value: V,
    height: i32,
    left: Option<Box<AVLTreeNode<K, V>>>,
    right: Option<Box<AVLTreeNode<K, V>>>,
}

impl<K, V> AVLTreeNode<K, V> {
    fn new(key: K, value: V) -> Self {
        AVLTreeNode {
            key,
            value,
            height: 1,
            left: None,
            right: None,
        }
    }
}

#[derive(Debug)]
struct AVLTree<K, V> {
    root: Option<Box<AVLTreeNode<K, V>>>,
    size: usize,
}

impl<K: Ord, V> AVLTree<K, V> {
    fn new() -> Self {
        AVLTree {
            root: None,
            size: 0,
        }
    }

    fn insert(&mut self, key: K, value: V) {
        self.root = Self::insert_node(self.root.take(), key, value);
        self.size += 1;
    }

    fn insert_node(
        node: Option<Box<AVLTreeNode<K, V>>>,
        key: K,
        value: V,
    ) -> Option<Box<AVLTreeNode<K, V>>> {
        if let Some(mut node) = node {
            match key.cmp(&node.key) {
                Ordering::Less => {
                    node.left = Self::insert_node(node.left.take(), key, value);
                }
                Ordering::Greater => {
                    node.right = Self::insert_node(node.right.take(), key, value);
                }
                Ordering::Equal => {
                    node.value = value;
                }
            }
            node.height = 1 + std::cmp::max(Self::height(&node.left), Self::height(&node.right));
            Some(Self::balance(node))
        } else {
            Some(Box::new(AVLTreeNode::new(key, value)))
        }
    }

    fn height(node: &Option<Box<AVLTreeNode<K, V>>>) -> i32 {
        node.as_ref().map_or(0, |n| n.height)
    }

    fn balance_factor(node: &Box<AVLTreeNode<K, V>>) -> i32 {
        Self::height(&node.left) - Self::height(&node.right)
    }

    fn balance(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let balance_factor = Self::balance_factor(&node);
        if balance_factor > 1 {
            if Self::balance_factor(node.left.as_ref().unwrap()) < 0 {
                node.left = Some(Self::rotate_left(node.left.take().unwrap()));
            }
            return Self::rotate_right(node);
        }
        if balance_factor < -1 {
            if Self::balance_factor(node.right.as_ref().unwrap()) > 0 {
                node.right = Some(Self::rotate_right(node.right.take().unwrap()));
            }
            return Self::rotate_left(node);
        }
        node
    }

    fn rotate_left(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let mut new_root = node.right.take().unwrap();
        node.right = new_root.left.take();
        new_root.left = Some(node);
        new_root.height = 1 + std::cmp::max(Self::height(&new_root.left), Self::height(&new_root.right));
        new_root.left.as_mut().unwrap().height = 1 + std::cmp::max(Self::height(&new_root.left.as_ref().unwrap().left), Self::height(&new_root.left.as_ref().unwrap().right));
        new_root
    }

    fn rotate_right(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let mut new_root = node.left.take().unwrap();
        node.left = new_root.right.take();
        new_root.right = Some(node);
        new_root.height = 1 + std::cmp::max(Self::height(&new_root.left), Self::height(&new_root.right));
        new_root.right.as_mut().unwrap().height = 1 + std::cmp::max(Self::height(&new_root.right.as_ref().unwrap().left), Self::height(&new_root.right.as_ref().unwrap().right));
        new_root
    }

    fn num_entries(&self) -> usize {
        self.size
    }

    fn validate(&self) {
        let mut counter = -1;
        Self::validate_subtree(&self.root, &mut counter);
    }

    fn validate_subtree(node: &Option<Box<AVLTreeNode<K, V>>>, counter: &mut i32) -> i32 {
        if let Some(node) = node {
            let left_height = Self::validate_subtree(&node.left, counter);
            assert!(*counter < node.key as i32);
            *counter = node.key as i32;
            let right_height = Self::validate_subtree(&node.right, counter);
            assert!(left_height - right_height < 2 && right_height - left_height < 2);
            std::cmp::max(left_height, right_height) + 1
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_avl_tree_new() {
        let tree = AVLTree::<i32, i32>::new();
        assert_eq!(tree.root, None);
        assert_eq!(tree.num_entries(), 0);
    }

    #[test]
    fn test_avl_tree_insert_lookup() {
        let mut tree = AVLTree::<i32, i32>::new();
        for i in 0..NUM_TEST_VALUES as i32 {
            tree.insert(i, i);
            assert_eq!(tree.num_entries(), (i + 1) as usize);
            tree.validate();
        }
    }
}
